#!/usr/bin/python
from pylab import *
from scipy import linalg
from scipy.ndimage.morphology import binary_erosion
from scipy.ndimage import interpolation,filters
from ocrolib import dbhelper,improc
import numpy,pylab,random,sqlite3,cv,random,signal,codecs,os,re,glob,sys
from optparse import OptionParser
import collections
from sutils import *
from ocrolib import docproc
from pdb import pm
import codecs

def tess_readchars(fname,boxfile=None,d=1,pad=2):
    """Read characters in Tesseract training format."""
    os.system("convert '%s' /tmp/image.png"%fname)
    image = imread("/tmp/image.png")
    h,w = image.shape
    if boxfile is None:
        boxfile = re.sub(r'(\.g4)*\.tiff?','.box',fname)
    with codecs.open(boxfile,"r","utf8") as stream:
        lineno = 0
        for line in stream.readlines():
            f = line.split()
            c = f[0]
            x0,y0,x1,y1 = [int(x) for x in f[1:5]]
            cimage = image[h-y1-1-d:h-y0-1+d,x0-d:x1+1].copy()
            # if "georgiai" in fname and lineno==1682: raise Exception("debug")
            cimage /= amax(cimage)
            cimage = improc.pad_by(1-cimage,pad)
            yield c,cimage
            lineno += 1

def tess_allchars(dir="tessdata/eng"):
    """Read a directory of Tesseract training images."""
    for fname in glob.glob(dir+"/*.g4.tif"):
        fid = re.sub(r'(\.g4)*\.tif$','',re.sub(r'^.*/','',fname))
        print "opening",fname,fid
        for c,img in readchars(fname):
            yield record(cls=c,image=img,font=fid)
            
import argparse
parser = argparse.ArgumentParser(description = "Put boxfile training data into a patches file.")
parser.add_argument('inputs',nargs="*")
parser.add_argument('-o','--output',default="tess.h5",help="output file")
parser.add_argument('-d','--dimension',default=32,type=int,help="character dimension")
parser.add_argument('-s','--split',default="all",help="all, test, or train")
args = parser.parse_args()
# args = parser.parse_args(["tessdata/eng/eng.georgiai.g4.tif"])
if args.inputs==[]: parser.print_help(); sys.exit(0)

d = args.dimension

count = 0
with tables.openFile(args.output,"w") as odb:
    patches = odb.createEArray(odb.root,'patches',tables.Float32Atom(),shape=(0,d,d),filters=tables.Filters(9))
    classes = odb.createEArray(odb.root,'classes',tables.Int64Atom(),shape=(0,),filters=tables.Filters(9))
    for fname in args.inputs:
        print "===",fname
        lineno = 0
        for c,cimage in tess_readchars(fname):
            if args.split!="all":
                if (args.split=="test" and count%10!=9) or (args.split=="train" and count%10==9):
                    continue
            cimage = docproc.isotropic_rescale(cimage,d)
            cimage = csnormalize(cimage)
            # if len(odb.root.classes)==91106: print fname,count,c,lineno
            patches.append([cimage])
            classes.append([uencode(c)])
            count += 1
            lineno += 1
